/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.zookeeper.server.quorum;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.nio.channels.UnresolvedAddressException;
import java.util.Enumeration;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.NoSuchElementException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.Date;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 每台服务器在启动的过程中，会启动一个QuorumPeerManager，负责各台服务器之间的底层Leader选举过程中的网络通信。
 *
 * 内部类
 *   SendWorker类作为网络IO的发送者，每个SendWorker不断的从对应的消息发送队列中获取一个消息来发送，并将这个消息放入lastMessageSent中，如果队列为空，则从lastMessageSent取出最后一个消息重新发送，可解决接收方没有正确接收或处理消息的问题
 *   RecvWorker类作为网络IO的接收者，每个RecvWorker不断的从这个TCP连接中读取消息，并将其保存到recvQueue队列中。
 *   Message类定义了消息结构，包含sid以及消息体ByteBuffer
 *   Listener类作为electionPort端口的监听器，等待其他机器的连接
 */

public class QuorumCnxManager {
    private static final Logger LOG = LoggerFactory.getLogger(QuorumCnxManager.class);

    /** 接收队列的长度 */
    static final int RECV_CAPACITY = 100;
    /** 发送队列的长度  */
    static final int SEND_CAPACITY = 1;

    static final int PACKETMAXSIZE = 1024 * 1024;

    /** Maximum number of attempts to connect to a peer */
    static final int MAX_CONNECTION_ATTEMPTS = 2;
    
    /** Negative counter for observer server ids. */
    private long observerCounter = -1;
    
    /** Connection time out value in milliseconds  */
    private int cnxTO = 5000;
    
    /** Local IP address */
    final QuorumPeer self;

    /** SendWorker是消息发送器，这是按照SID分组的消息发送器集合 */
    final ConcurrentHashMap<Long, SendWorker> senderWorkerMap;
    /** 消息发送队列，按照SID分组，用于保存待发送的消息，从而保证了各台机器之间的消息发送互不影响 */
    final ConcurrentHashMap<Long, ArrayBlockingQueue<ByteBuffer>> queueSendMap;
    /** 最近发送过的消息，为每个SID保留最近发送过的消息 */
    final ConcurrentHashMap<Long, ByteBuffer> lastMessageSent;

    /** 消息接收队列，用于存放从其他服务器接收到的消息 */
    public final ArrayBlockingQueue<Message> recvQueue;
    /** Object to synchronize access to recvQueue */
    private final Object recvQLock = new Object();

    /** Shutdown flag */
    volatile boolean shutdown = false;

    /** Listener thread */
    public final Listener listener;

    /** Counter to count worker threads */
    private AtomicInteger threadCnt = new AtomicInteger(0);

    static public class Message {
        Message(ByteBuffer buffer, long sid) {
            this.buffer = buffer;
            this.sid = sid;
        }

        ByteBuffer buffer;
        long sid;
    }

    public QuorumCnxManager(QuorumPeer self) {
        this.recvQueue = new ArrayBlockingQueue<Message>(RECV_CAPACITY);
        this.queueSendMap = new ConcurrentHashMap<Long, ArrayBlockingQueue<ByteBuffer>>();
        this.senderWorkerMap = new ConcurrentHashMap<Long, SendWorker>();
        this.lastMessageSent = new ConcurrentHashMap<Long, ByteBuffer>();
        
        String cnxToValue = System.getProperty("zookeeper.cnxTimeout");
        if(cnxToValue != null){
            this.cnxTO = new Integer(cnxToValue); 
        }
        
        this.self = self;

        // Starts listener thread that waits for connection requests 
        listener = new Listener();
    }

    /**
     * Invokes initiateConnection for testing purposes
     * 
     * @param sid
     */
    public void testInitiateConnection(long sid) throws Exception {
        Socket sock = new Socket();
        setSockOpts(sock);
        sock.connect(self.getVotingView().get(sid).electionAddr, cnxTO);
        initiateConnection(sock, sid);
    }
    
    /**
     *  初始化连接。 在两两创建连接时，有个规则：
     *  只允许SID大的服务器主动和其他服务器建立连接，否则断开连接。在receiveConnection方法中，服务器会接受远程SID比自己大的连接。从而避免了两台服务器之间的重复连接。
     */
    public boolean initiateConnection(Socket sock, Long sid) {
        DataOutputStream dout = null;
        try {
            // 直接用java的BIO架构，进行节点之间的通信
            dout = new DataOutputStream(sock.getOutputStream());
            dout.writeLong(self.getId()); // 发送本机sid
            dout.flush();
        } catch (IOException e) {
            LOG.warn("Ignoring exception reading or writing challenge: ", e);
            closeSocket(sock);
            return false;
        }

        // 发送连接的时候，只让大sid给小sid发送，如果当前sid小，那就close掉
        if (sid > self.getId()) {
            LOG.info("Have smaller server identifier, so dropping the " + "connection: (" + sid + ", " + self.getId() + ")");
            closeSocket(sock);
        } else {
            // 自己sid大，初始化SendWorker和RecvWorker
            SendWorker sw = new SendWorker(sock, sid);
            RecvWorker rw = new RecvWorker(sock, sid, sw); // rw记录sw
            // sw记录rw
            sw.setRecv(rw);

            SendWorker vsw = senderWorkerMap.get(sid);
            
            if(vsw != null) {
                // finish掉sid对应的SendWorker,vsw
                vsw.finish();
            }

            // 放入新的SendWorker,sw
            senderWorkerMap.put(sid, sw);
            if (!queueSendMap.containsKey(sid)) {
                queueSendMap.put(sid, new ArrayBlockingQueue<ByteBuffer>( SEND_CAPACITY));
            }
            
            sw.start();
            rw.start();
            return true;
        }
        return false;
    }

    
    
    /**
     * 接收connect
     * 
     */
    public boolean receiveConnection(Socket sock) {
        Long sid = null;
        
        try {
            // Read server id
            DataInputStream din = new DataInputStream(sock.getInputStream());
            sid = din.readLong();
            if (sid == QuorumPeer.OBSERVER_ID) {
                sid = observerCounter--;
                LOG.info("Setting arbitrary identifier to observer: " + sid);
            }
        } catch (IOException e) {
            closeSocket(sock);
            LOG.warn("Exception reading or writing challenge: " + e.toString());
            return false;
        }
        // 如果自己id大，就close掉当前连接(当前是小sid发给大sid的连接)，自己再去连对方sid
        if (sid < self.getId()) {
            SendWorker sw = senderWorkerMap.get(sid);
            if (sw != null) {
                sw.finish();
            }

            LOG.debug("Create new connection to server: " + sid);
            closeSocket(sock);
            connectOne(sid);

        } else {
            SendWorker sw = new SendWorker(sock, sid);
            RecvWorker rw = new RecvWorker(sock, sid, sw);
            sw.setRecv(rw);

            SendWorker vsw = senderWorkerMap.get(sid);
            
            if(vsw != null) {
                vsw.finish();
            }
            
            senderWorkerMap.put(sid, sw);
            
            if (!queueSendMap.containsKey(sid)) {
                queueSendMap.put(sid, new ArrayBlockingQueue<ByteBuffer>( SEND_CAPACITY));
            }
            
            sw.start();
            rw.start();
            return true;
        }
        return false;
    }

    /**
     * 将消息根据sid添加进recv队列或者send队列,间接调用send，recv的生产
     */
    public void toSend(Long sid, ByteBuffer b) {
        // 如果发送给自己,加入recv队列
        if (self.getId() == sid) {
             b.position(0);
             addToRecvQueue(new Message(b.duplicate(), sid));
        }
        // 如果要发送给其他机器, 首先判断本地 消息发送队列中是否存在 sid的消息集合，
        // 没有就new一个，有的话，就放到该消息集合里面
        else {
            // 如果发送map没有记录这个sid
             if (!queueSendMap.containsKey(sid)) {
                 // 阻塞队列长度为1
                 ArrayBlockingQueue<ByteBuffer> bq = new ArrayBlockingQueue<>(SEND_CAPACITY);
                 queueSendMap.put(sid, bq);
                 addToSendQueue(bq, b);
             } else {
                 ArrayBlockingQueue<ByteBuffer> bq = queueSendMap.get(sid);
                 if(bq != null){
                     addToSendQueue(bq, b);
                 } else {
                     LOG.error("No queue for server " + sid);
                 }
             }
            // 和这个sid建立连接
             connectOne(sid);
        }
    }
    
    /**
     * 连接上某个sid的server
     *  @param sid  server id
     */
    synchronized void connectOne(long sid){
        // 只有senderWorkerMap中没有该sid，才会执行
        if (senderWorkerMap.get(sid) == null){
            InetSocketAddress electionAddr;
            if (self.quorumPeers.containsKey(sid)) {
                // 从配置文件获取对应sid机器的选举端口
                electionAddr = self.quorumPeers.get(sid).electionAddr;
            } else {
                LOG.warn("Invalid server id: " + sid);
                return;
            }
            try {
                // zookeeper底层，是直接基于Java Socket进行TCP协议通信的。
                Socket sock = new Socket();
                setSockOpts(sock);
                // 连接上对应socket
                sock.connect(self.getView().get(sid).electionAddr, cnxTO);
                // 核心方法，初始化连接， 为该节点新建SendWork 放入senderWorkerMap中
                initiateConnection(sock, sid);
            } catch (UnresolvedAddressException e) {
                LOG.warn("Cannot open channel to " + sid + " at election address " + electionAddr, e);
                throw e;
            } catch (IOException e) {
                LOG.warn("Cannot open channel to " + sid + " at election address " + electionAddr, e);
            }
        } else {
            LOG.debug("There is a connection already for server " + sid);
        }
    }
    
    
    /**
     * 连接所有queueSendMap的sid
     */
    public void connectAll(){
        // 把所有需要发送消息的机器sid都连接上
        long sid;
        // 连接所有queueSendMap记录的sid
        for(Enumeration<Long> en = queueSendMap.keys(); en.hasMoreElements();){
            sid = en.nextElement();
            connectOne(sid);
        }      
    }
    

    /**
     * 是否发送过消息
     */
    boolean haveDelivered() {
        // 如果有一个队列是空的，代表发送过了
        for (ArrayBlockingQueue<ByteBuffer> queue : queueSendMap.values()) {
            if (queue.size() == 0) {
                return true;
            }
        }
        return false;
    }

    /**
     * Flag that it is time to wrap up all activities and interrupt the listener.
     */
    public void halt() {
        shutdown = true;
        LOG.debug("Halting listener");
        listener.halt();
        
        softHalt();
    }
   
    /**
     * A soft halt simply finishes workers.
     */
    public void softHalt() {
        for (SendWorker sw : senderWorkerMap.values()) {
            LOG.debug("Halting sender: " + sw);
            sw.finish();
        }
    }

    /**
     * Helper method to set socket options.
     * 
     * @param sock Reference to socket
     */
    private void setSockOpts(Socket sock) throws SocketException {
        sock.setTcpNoDelay(true);
        sock.setSoTimeout(self.tickTime * self.syncLimit);
    }

    /**
     * Helper method to close a socket.
     * 
     * @param sock Reference to socket
     */
    private void closeSocket(Socket sock) {
        try {
            sock.close();
        } catch (IOException ie) {
            LOG.error("Exception while closing", ie);
        }
    }

    /**
     * Return number of worker threads
     */
    public long getThreadCount() {
        return threadCnt.get();
    }
    /**
     * Return reference to QuorumPeer
     */
    public QuorumPeer getQuorumPeer() {
        return self;
    }

    /**
     * 通过 ServerSocket连接集群某个id对应的ip和端口号。
     */
    public class Listener extends Thread {

        volatile ServerSocket ss = null;

        /**
         * Sleeps on accept().
         */
        @Override
        public void run() {
            int numRetries = 0;
            while((!shutdown) && (numRetries < 3)){
                try {
                    ss = new ServerSocket();
                    ss.setReuseAddress(true);
                    int port = self.quorumPeers.get(self.getId()).electionAddr.getPort();
                    InetSocketAddress addr = new InetSocketAddress(port);
                    LOG.info("My election bind port: " + addr.toString());
                    setName(self.quorumPeers.get(self.getId()).electionAddr.toString());
                    ss.bind(addr);
                    while (!shutdown) {
                        Socket client = ss.accept();
                        setSockOpts(client);
                        LOG.info("Received connection request " + client.getRemoteSocketAddress());
                        receiveConnection(client);
                        numRetries = 0;
                    }
                } catch (IOException e) {
                    LOG.error("Exception while listening", e);
                    numRetries++;
                    try {
                        ss.close();
                        Thread.sleep(1000);
                    } catch (IOException ie) {
                        LOG.error("Error closing server socket", ie);
                    } catch (InterruptedException ie) {
                        LOG.error("Interrupted while sleeping. " + "Ignoring exception", ie);
                    }
                }
            }
            LOG.info("Leaving listener");
            if (!shutdown) {
                LOG.error("As I'm leaving the listener thread, "
                        + "I won't be able to participate in leader "
                        + "election any longer: "
                        + self.quorumPeers.get(self.getId()).electionAddr);
            }
        }
        
        /**
         * Halts this listener thread.
         */
        void halt(){
            try{
                LOG.debug("Trying to close listener: " + ss);
                if(ss != null) {
                    LOG.debug("Closing listener: " + self.getId());
                    ss.close();
                }
            } catch (IOException e){
                LOG.warn("Exception when shutting down listener: " + e);
            }
        }
    }

    /**
     * 这个类作为“发送者”，继承ZooKeeperThread，线程不断地从发送队列取出，发给对应sid的机器.
     *
     * 在SendWorker中，一旦Zookeeper发现针对当前服务器的消息发送队列为空，那么此时需要从lastMessageSent中取出一个最近发送过的消息来进行再次发送，
     * 这是为了解决接收方在消息接收前或者接收到消息后服务器挂了，导致消息尚未被正确处理。
     * 同时，Zookeeper能够保证接收方在处理消息时，会对重复消息进行正确的处理。
     */
    class SendWorker extends Thread {
        Long sid;  // 目标机器sid，不是当前机器sid
        Socket sock;
        RecvWorker recvWorker; // 该sid对应的RecvWorker
        volatile boolean running = true;
        DataOutputStream dout;

        /**
         * SendWorker 线程 有目标机器的 sid号，会不停的从 queueSendMap 根据sid拿到对应的发送消息队列，
         * 从该消息队列里拉取消息，通过SendWorker线程中的DataOutputStream发送出去。
         * @param sock Socket to remote peer
         * @param sid Server identifier of remote peer
         */
        SendWorker(Socket sock, Long sid) {
            super("SendWorker:" + sid);
            this.sid = sid;
            this.sock = sock;
            recvWorker = null;
            try {
                dout = new DataOutputStream(sock.getOutputStream());
            } catch (IOException e) {
                LOG.error("Unable to access socket output stream", e);
                closeSocket(sock);
                running = false;
            }
            LOG.debug("Address of remote peer: " + this.sid);
        }

        synchronized void setRecv(RecvWorker recvWorker) {
            this.recvWorker = recvWorker;
        }

        /**
         * Returns RecvWorker that pairs up with this SendWorker.
         * @return RecvWorker
         */
        synchronized RecvWorker getRecvWorker(){
            return recvWorker;
        }
                
        synchronized boolean finish() {
            if(!running){
                return running;
            }
            
            running = false;
            closeSocket(sock);

            this.interrupt();
            if (recvWorker != null) {
                recvWorker.finish();
            }

            senderWorkerMap.remove(sid, this);
            threadCnt.decrementAndGet();
            return running;
        }
        
        synchronized void send(ByteBuffer b) throws IOException {
            byte[] msgBytes = new byte[b.capacity()];
            try {
                b.position(0);
                b.get(msgBytes);
            } catch (BufferUnderflowException be) {
                LOG.error("BufferUnderflowException ", be);
                return;
            }
            dout.writeInt(b.capacity());
            dout.write(b.array());
            dout.flush();
        }

        @Override
        public void run() {
            threadCnt.incrementAndGet(); // 线程数+1
            try {
                // 找到sid对应需要send的队列
                ArrayBlockingQueue<ByteBuffer> bq = queueSendMap.get(sid);
                if (bq == null || isSendQueueEmpty(bq)) {
                    // 如果没有什么发的，就把上一次发的再发一遍(重发能够正确处理)
                   ByteBuffer b = lastMessageSent.get(sid);
                   if (b != null) {
                       send(b); //发送
                   }
                }
            } catch (IOException e) {
                LOG.error("Failed to send last message. Shutting down thread.", e);
                this.finish();
            }
            
            try {
                while (running && !shutdown && sock != null) {
                    ByteBuffer b = null;
                    try {
                        ArrayBlockingQueue<ByteBuffer> bq = queueSendMap.get(sid);
                        if (bq != null) {
                            // 从发送队列里面取出消息
                            b = pollSendQueue(bq, 1000, TimeUnit.MILLISECONDS);
                        } else {
                            LOG.error("No queue of incoming messages for " + "server " + sid);
                            break;
                        }

                        if(b != null){
                            // 更新最后一次发送的
                            lastMessageSent.put(sid, b);
                            send(b); // 发送
                        }
                    } catch (InterruptedException e) {
                        LOG.warn("Interrupted while waiting for message on queue", e);
                    }
                }
            } catch (Exception e) {
                LOG.warn("Exception when using channel: for id " + sid + " my id = " +  self.getId() + " error = " + e);
            }
            this.finish();
            LOG.warn("Send worker leaving thread");
        }
    }

    /**
     * 这个类作为“接受者”，类似SendWorker，继承ZooKeeperThread，线程不断地从网络IO中读取数据，放入recvQueue接收队列
     */
    class RecvWorker extends Thread {
        Long sid; //来源方sid
        Socket sock;
        volatile boolean running = true;
        DataInputStream din;
        final SendWorker sw;

        RecvWorker(Socket sock, Long sid, SendWorker sw) {
            super("RecvWorker:" + sid);
            this.sid = sid;
            this.sock = sock;
            this.sw = sw;
            try {
                din = new DataInputStream(sock.getInputStream());
                sock.setSoTimeout(0);
            } catch (IOException e) {
                LOG.error("Error while accessing socket for " + sid, e);
                closeSocket(sock);
                running = false;
            }
        }
        
        /**
         * Shuts down this worker
         * @return boolean  Value of variable running
         */
        synchronized boolean finish() {
            if(!running){
                return running;
            }
            running = false;            

            this.interrupt();
            threadCnt.decrementAndGet();
            return running;
        }

        @Override
        public void run() {
            threadCnt.incrementAndGet();
            try {
                while (running && !shutdown && sock != null) {
                    // 获取长度
                    int length = din.readInt();
                    if (length <= 0 || length > PACKETMAXSIZE) {
                        throw new IOException( "Received packet with invalid packet: " + length);
                    }
                    byte[] msgArray = new byte[length];
                    din.readFully(msgArray, 0, length);
                    // 解析出byteBuffer
                    ByteBuffer message = ByteBuffer.wrap(msgArray);
                    // 加入接收队列
                    addToRecvQueue(new Message(message.duplicate(), sid));
                }
            } catch (Exception e) {
                LOG.warn("Connection broken for id " + sid + ", my id = " +  self.getId() + ", error = " , e);
            } finally {
                LOG.warn("Interrupting SendWorker");
                sw.finish();
                if (sock != null) {
                    closeSocket(sock);
                }
            }
        }
    }

    /**
     * 每个sender的队列长度都是1,为了避免发送旧的数据，因此会把旧的remove掉
     *
     * @param queue Reference to the Queue
     * @param buffer Reference to the buffer to be inserted in the queue
     */
    private void addToSendQueue(ArrayBlockingQueue<ByteBuffer> queue, ByteBuffer buffer) {
        // 发送队列长度为1，如果满了就remove，然后add
        if (queue.remainingCapacity() == 0) {
            try {
                queue.remove();
            } catch (NoSuchElementException ne) {
                LOG.debug("Trying to remove from an empty " + "Queue. Ignoring exception " + ne);
            }
        }
        try {
            queue.add(buffer);
        } catch (IllegalStateException ie) {
            LOG.error("Unable to insert an element in the queue " + ie);
        }
    }

    /**
     * Returns true if queue is empty.
     * @param queue Reference to the queue
     * @return true if the specified queue is empty
     */
    private boolean isSendQueueEmpty(ArrayBlockingQueue<ByteBuffer> queue) {
        return queue.isEmpty();
    }

    /**
     * Retrieves and removes buffer at the head of this queue,
     * waiting up to the specified wait time if necessary for an element to
     * become available.
     *
     * {@link ArrayBlockingQueue#poll(long, java.util.concurrent.TimeUnit)}
     */
    private ByteBuffer pollSendQueue(ArrayBlockingQueue<ByteBuffer> queue, long timeout, TimeUnit unit) throws InterruptedException {
       return queue.poll(timeout, unit);
    }

    /**
     * Inserts an element in the {@link #recvQueue}. If the Queue is full, this
     * methods removes an element from the head of the Queue and then inserts
     * the element at the tail of the queue.
     *
     * This method is synchronized to achieve fairness between two threads that
     * are trying to insert an element in the queue. Each thread checks if the
     * queue is full, then removes the element at the head of the queue, and
     * then inserts an element at the tail. This three-step process is done to
     * prevent a thread from blocking while inserting an element in the queue.
     * If we do not synchronize the call to this method, then a thread can grab
     * a slot in the queue created by the second thread. This can cause the call
     * to insert by the second thread to fail.
     * Note that synchronizing this method does not block another thread
     * from polling the queue since that synchronization is provided by the
     * queue itself.
     *
     * @param msg Reference to the message to be inserted in the queue
     */
    public void addToRecvQueue(Message msg) {
        synchronized(recvQLock) {
            if (recvQueue.remainingCapacity() == 0) {
                try {
                    recvQueue.remove();
                } catch (NoSuchElementException ne) {
                     LOG.debug("Trying to remove from an empty " + "recvQueue. Ignoring exception " + ne);
                }
            }
            try {
                recvQueue.add(msg);
            } catch (IllegalStateException ie) {
                LOG.error("Unable to insert element in the recvQueue " + ie);
            }
        }
    }

    /**
     * Retrieves and removes a message at the head of this queue,
     * waiting up to the specified wait time if necessary for an element to
     * become available.
     *
     * {@link ArrayBlockingQueue#poll(long, java.util.concurrent.TimeUnit)}
     */
    public Message pollRecvQueue(long timeout, TimeUnit unit)
       throws InterruptedException {
       return recvQueue.poll(timeout, unit);
    }
}
